/*
 * Copyright 2023 lzhpo
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.ruoyi.web.controller.chat;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.lzhpo.chatgpt.OpenAiClient;
import com.lzhpo.chatgpt.OpenAiKeyWrapper;
import com.lzhpo.chatgpt.OpenAiUrl;
import com.lzhpo.chatgpt.entity.chat.ChatCompletionMessage;
import com.lzhpo.chatgpt.entity.chat.ChatCompletionRequest;
import com.lzhpo.chatgpt.entity.chat.ChatCompletionResponse;
import com.lzhpo.chatgpt.entity.embeddings.EmbeddingRequest;
import com.lzhpo.chatgpt.entity.embeddings.EmbeddingResponse;
import com.lzhpo.chatgpt.sse.SseEventSourceListener;
import com.ruoyi.domain.GoogleSearchResponse;
import com.ruoyi.service.IChatService;
import com.ruoyi.service.MicroSoftApiClient;
import com.ruoyi.web.controller.chat.util.ChangeDTO;
import com.ruoyi.web.controller.chat.util.ChangeUtil;
import com.ruoyi.web.controller.chat.util.ChatUtil;
import com.ruoyi.web.core.config.ChatModel;
import com.ruoyi.web.core.config.HfThreadLocal;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.ModelAndView;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.validation.Valid;
import javax.validation.constraints.NotBlank;
import java.net.URLEncoder;
import java.util.List;

import static com.lzhpo.chatgpt.OpenAiUrl.CHAT_COMPLETIONS;

/**
 * @author lzhpo
 */
@Slf4j
@Controller
@RequestMapping("/chat-api")
@RequiredArgsConstructor
@Api("OpenAi接口")
@Validated
public class OpenAiController extends OpenAiAuthController {

    @Autowired
    private IChatService chatService;

    private final OpenAiClient openAiClient;
    private final OpenAiKeyWrapper openAiKeyWrapper;

    @Autowired
    private MicroSoftApiClient microSoftApiClient;

    @ApiOperation("普通用户对话视图")
    @GetMapping("/page/chat")
    public ModelAndView chatView() {
        ModelAndView view = new ModelAndView("chat");
        view.addObject("apiKey", openAiKeyWrapper.wrap().next());
        return view;
    }

    @ApiOperation("sse对话视图")
    @GetMapping("/page/chat/sse")
    public ModelAndView sseChatView() {
        return new ModelAndView("sse-stream-chat");
    }

    @ApiOperation("websocket对话视图")
    @GetMapping("/page/chat/websocket")
    public ModelAndView websocketChatView() {
        return new ModelAndView("websocket-stream-chat");
    }

    @ApiOperation("普通chat接口")
    @ResponseBody
    @PostMapping("/chat")
    public ChatCompletionResponse chatCompletions(@Validated @RequestBody ChatCompletionRequest request) {
        String openAiKey = getOpenAiKey();
        ChatCompletionResponse chatCompletionResponse = null;

        String realModel = request.getModel();
        ChangeDTO change = ChangeUtil.change(realModel, openAiKey);
        HfThreadLocal.setModel(change);
        ChangeUtil.changeRequest(request, change);

        try {
            chatCompletionResponse = openAiClient.chatCompletions(request);
            ChangeUtil.changeResponse(chatCompletionResponse, change);
            return chatCompletionResponse;
        } finally {
            saveRecord(openAiKey, request, chatCompletionResponse);
            HfThreadLocal.clear();
        }
    }

    @ApiOperation("embeddings接口")
    @ResponseBody
    @PostMapping("/embeddings")
    public EmbeddingResponse embeddings(@Validated @RequestBody EmbeddingRequest request) {
        String openAiKey = getOpenAiKey();
        EmbeddingResponse embeddingResponse = null;

        String realModel = request.getModel();
        ChangeDTO change = ChangeUtil.change(realModel, openAiKey);
        HfThreadLocal.setModel(change);

        try {
            embeddingResponse = openAiClient.embeddings(request);
            return embeddingResponse;
        } finally {
            saveEmbeddingRecord(openAiKey, request, embeddingResponse);
            HfThreadLocal.clear();
        }
    }

    @ApiOperation("chat流式接口")
    @ResponseBody
    @GetMapping("/chat/stream")
    public HfSseEmitter streamChat(@NotBlank @RequestParam String request) {
        String openAiKey = getOpenAiKey();
        HfSseEmitter sseEmitter = new HfSseEmitter();

        ChatCompletionRequest chatCompletionRequest = JSONObject.parseObject(request, ChatCompletionRequest.class);

        String realModel = chatCompletionRequest.getModel();
        ChangeDTO change = ChangeUtil.change(realModel, openAiKey);
        HfThreadLocal.setModel(change);

        ChangeUtil.changeRequest(chatCompletionRequest, change);

        openAiClient.streamChatCompletions(chatCompletionRequest, new SseEventSourceListener(sseEmitter));
        saveRecordStream(openAiKey, chatCompletionRequest);
        return sseEmitter;
    }

    @ApiOperation("普通chat接口")
    @ResponseBody
    @PostMapping("/chatWithSearch")
    public ChatCompletionResponse chatCompletionsWithSearch(@Validated @RequestBody ChatCompletionRequest request) {
        String openAiKey = getOpenAiKey();
        //从message中获取最后一个user的问题
        String userQa = ChatUtil.getUserQa(request.getMessages());
        if (StringUtils.isNotBlank(userQa)) {
            String s = chatService.searchResult(userQa);
            if (StringUtils.isNotBlank(s)) {
                ChatUtil.appendSystem(request.getMessages(), s);
            }
        }

        return openAiClient.chatCompletions(request);
    }

    @ApiOperation("sse chat接口")
    @ResponseBody
    @GetMapping("/chat/sse")
    public SseEmitter sseStreamChat(@NotBlank @RequestParam String message) {
        String openAiKey = getOpenAiKey();
        SseEmitter sseEmitter = new SseEmitter();
        ChatCompletionRequest request = ChatCompletionRequest.create(message);
        openAiClient.streamChatCompletions(request, new SseEventSourceListener(sseEmitter));
        return sseEmitter;
    }

    @ApiOperation("私人语料")
    @ResponseBody
    @PostMapping("/chatWithInner")
    public ChatCompletionResponse chatCompletionsWithInner(@Validated @RequestBody ChatCompletionRequest request) {
        String openAiKey = getOpenAiKey();
        //从message中获取最后一个user的问题
        String userQa = ChatUtil.getUserQa(request.getMessages());
        if (StringUtils.isNotBlank(userQa)) {
            String s = chatService.innerChat(userQa);
            if (StringUtils.isNotBlank(s)) {
                ChatUtil.appendSystem(request.getMessages(), s);
            }
        }

        System.out.println(request);
        return openAiClient.chatCompletions(request);
    }

    public static void main(String[] args) {
        System.out.println(URLEncoder.encode("AN: 测试你的游戏上瘾指数, AN_detail: 游戏上瘾是指对游戏内容的过度沉迷而影响到了日常生活和工作，那么你的游戏上瘾指数到达了什么级别呢？下面的测试题将告诉你。, QA: [  q: 1.你玩游戏时，通常会花费多长时间？, a: [  name: 半个小时以内, value: 0 ,  name: 1 - 2小时, value: 1 ,  name: 3 - 4小时, value: 2 ,  name: 5小时及以上, value: 3  ]  ], R: [  r: 3, t: 轻度瘾者, s: 100分, c: 虽然你对游戏有些爱好，但并没有到达严重上瘾的程度，可以在适当的时间内保持适度的游戏时间，并且不影响正常生活和工作。  ]  按照以上json格式，写出和【你是不解风情的爱情木头吗】相关的心理测试题。要求6个问题，AN为【你是不解风情的爱情木头吗】，AN_detail要40个字以上，value从0开始计数，s给出评分总分100分，c控制在100个字以内"));
    }
}
